{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.debugger import set_trace\n",
    "import numpy as np\n",
    "import pprint\n",
    "import sys\n",
    "if \"../\" not in sys.path:\n",
    "  sys.path.append(\"../\") \n",
    "from lib.envs.gridworld import GridworldEnv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "pp = pprint.PrettyPrinter(indent=2)\n",
    "env = GridworldEnv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def policy_eval(policy, env, discount_factor=1.0, theta=0.00001):\n",
    "    \"\"\"\n",
    "    Evaluate a policy given an environment and a full description of the environment's dynamics.\n",
    "    \n",
    "    Args:\n",
    "        policy: [S, A] shaped matrix representing the policy.\n",
    "        env: OpenAI env. env.P represents the transition probabilities of the environment.\n",
    "            env.P[s][a] is a list of transition tuples (prob, next_state, reward, done).\n",
    "            env.nS is a number of states in the environment. \n",
    "            env.nA is a number of actions in the environment.\n",
    "        theta: We stop evaluation once our value function change is less than theta for all states.\n",
    "        discount_factor: Gamma discount factor.\n",
    "    \n",
    "    Returns:\n",
    "        Vector of length env.nS representing the value function.\n",
    "    \"\"\"\n",
    "    # Start with a random (all 0) value function\n",
    "    V = np.zeros(env.nS)\n",
    "    while True:\n",
    "        delta = 0\n",
    "        # For each state, perform a \"full backup\"\n",
    "        for s in range(env.nS):\n",
    "            v = 0\n",
    "            # Look at the possible next actions\n",
    "            for a, action_prob in enumerate(policy[s]):\n",
    "                # For each action, look at the possible next states...\n",
    "                for  prob, next_state, reward, done in env.P[s][a]:\n",
    "                    # Calculate the expected value. Ref: Sutton book eq. 4.6.\n",
    "                    v += action_prob * prob * (reward + discount_factor * V[next_state])\n",
    "            # How much our value function changed (across any states)\n",
    "            delta = max(delta, np.abs(v - V[s]))\n",
    "            V[s] = v\n",
    "        # Stop evaluating once our value function change is below a threshold\n",
    "        if delta < theta:\n",
    "            break\n",
    "    return np.array(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_policy = np.ones([env.nS, env.nA]) / env.nA\n",
    "v = policy_eval(random_policy, env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Value Function:\n",
      "[  0.         -13.99993529 -19.99990698 -21.99989761 -13.99993529\n",
      " -17.9999206  -19.99991379 -19.99991477 -19.99990698 -19.99991379\n",
      " -17.99992725 -13.99994569 -21.99989761 -19.99991477 -13.99994569\n",
      "   0.        ]\n",
      "\n",
      "Reshaped Grid Value Function:\n",
      "[[  0.         -13.99993529 -19.99990698 -21.99989761]\n",
      " [-13.99993529 -17.9999206  -19.99991379 -19.99991477]\n",
      " [-19.99990698 -19.99991379 -17.99992725 -13.99994569]\n",
      " [-21.99989761 -19.99991477 -13.99994569   0.        ]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"Value Function:\")\n",
    "print(v)\n",
    "print(\"\")\n",
    "\n",
    "print(\"Reshaped Grid Value Function:\")\n",
    "print(v.reshape(env.shape))\n",
    "print(\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test: Make sure the evaluated policy is what we expected\n",
    "expected_v = np.array([0, -14, -20, -22, -14, -18, -20, -20, -20, -20, -18, -14, -22, -20, -14, 0])\n",
    "np.testing.assert_array_almost_equal(v, expected_v, decimal=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    ""
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3.0
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}